import os
import hashlib
import sys

from Crypto.Cipher import AES
from Crypto.Util import Counter

import PBKDF


class KeyManager:

    def __init__(self, aesKey):
        self.key = self.__deriveKey(aesKey)
    
    def __deriveKey(self, key):
        if not os.path.exists("%s/.dropsecure" % os.getcwd()):
            os.mkdir("%s/.dropsecure" % os.getcwd())
        salt = ''
        if not os.path.exists("%s/.dropsecure/salt.aes" % os.getcwd()):
            salt = os.urandom(32)
            with open("%s/.dropsecure/salt.aes" % os.getcwd(), "wb") as f:
                f.write(salt)
        else:
            with open("%s/.dropsecure/salt.aes" % os.getcwd(), "rb") as f:
                salt = f.read()
        h = PBKDF.PBKDF2(key,salt)
        return h
    
    def getKey(self):
        return self.key
